# from generate_new_pic import generate_annotated_image
# Paddleocr目前支持的多语言语种可以通过修改lang参数进行切换
# 例如`ch`, `en`, `fr`, `german`, `korean`, `japan`


# 显示结果
# def show_result_as_png(img_path, txts, boxes, scores):
#     from PIL import Image
#     from paddleocr import draw_ocr
#     image = Image.open(img_path).convert('RGB')
#     im_show = draw_ocr(image, boxes, txts, scores, font_path='./fonts/simfang.ttf')
#     im_show = Image.fromarray(im_show)
#     im_show.save('fapiao2_result.jpg')


def image_ocr(parse_out_dir, image_saved_dir):
    """
    函数 image_ocr 读取解析目录中的图像文件信息，执行 OCR 操作并将结果保存。
    参数说明：
    - parse_out_dir: 解析后的输出目录，包含图像信息的 JSON 文件。
    - image_saved_dir: 图像文件保存的目录，用于 OCR 处理。
    """
    from paddleocr import PaddleOCR
    from preprocess.file_paser.paser_utils import traverse_dir, re_filter
    import json
    import os
    # need to run only once to download and load model into memory
    print("对图片进行OCR：")
    ocr = PaddleOCR(use_angle_cls=True, lang="ch")
    # 遍历指定目录下的所有 JSON 文件
    iter_file = traverse_dir(parse_out_dir, ['json'])
    exclude_content_set = {'CT)'}

    # 对于每个 JSON 文件，读取内容并执行 OCR 操作
    for json_path, _ in iter_file:

        # 打开 JSON 文件并加载数据
        with open(json_path, 'r', encoding='utf-8') as file:
            data_dict = json.load(file)

        # 遍历文件列表中的每个图像信息字典
        for dic in data_dict['file_list']:
            image_info = dic['image_info']  # 获取图像信息字典

            # 遍历每个图像项，执行 OCR 操作并更新内容
            for k, v in image_info.items():
                print('-' * 99)  # 打印分割线，便于查看控制台输出
                print(v['path'])  # 打印图像路径
                img_path = os.path.join(image_saved_dir, v['path'])  # 构建图像完整路径
                result = ocr.ocr(img_path, cls=True)
                if result and result[0]:
                    # 从图像中提取文本，进行正则过滤并去除空格
                    ocr_text = re_filter(' '.join([line[1][0] for line in result[0]])).strip()
                else:
                    ocr_text = ''
                if ocr_text in exclude_content_set:
                    image_info.pop(k)
                else:
                    image_info[k]['content'] = ocr_text

                    # 将处理后的数据重新写入 JSON 文件
        with open(json_path, 'w+', encoding='utf-8') as f:
            json.dump(data_dict, f, ensure_ascii=False)
